from typing import Any
from gremlin_python.structure.graph import GraphTraversalSource
from gremlin_python.process.graph_traversal import GraphTraversal


class QuerySet:
    """
    QuerySet is a wrapper around GraphTraversal object.
    """

    def __init__(self, g: GraphTraversal):
        """
        Initialize a QuerySet object with a GraphTraversal object.
        """
        assert g is not None

        self.source = g
        self.query = g.clone()

    ########## Delegate special methods to GraphTraversal ##########
    def __repr__(self):
        """
        See gremlin_python.process.traversal.py
        """
        return self.query.__repr__()

    def __eq__(self, other):
        """
        See gremlin_python.process.traversal.py
        """
        return self.query.__eq__(other.query)

    def __iter__(self):
        """
        See gremlin_python.process.traversal.py
        """
        return self.query.__iter__()

    def __next__(self):
        """
        See gremlin_python.process.traversal.py
        """
        return self.query.__next__()

    def __getitem__(self, index):
        """
        See gremlin_python.process.graph_traversal.py
        """
        return self.query.__getitem__(index)

    def __getattribute__(self, key: str):
        """
        See gremlin_python.process.graph_traversal.py
        """
        try:
            return object.__getattribute__(self, key)
        except AttributeError:
            return self.query.__getattribute__(key)

    ########## Delegate methods to GraphTraversal ##########
    def get(self):
        """
        If other methods are needed, use this to expose the
        original GraphTraversal object.
        """
        return self.query

    def apply(self, action, *args, **kwargs):
        """
        Apply an action to the QuerySet.
        This enables more flexible query building.
        All actions must at least take a GraphTraversal object as the first argument.
        """
        if action is not None:
            self.query = action(self.query, *args, **kwargs)
        return self

    ########## QuerySet methods ##########
    def clone(self):
        """
        Clone the QuerySet object.
        """
        return QuerySet(query=self.query.clone())
